# services/websocket_service.py

from fastapi import WebSocket, WebSocketDisconnect
from typing import List
from models import Message
from datetime import datetime

class ConnectionManager:
    """
    管理所有 WebSocket 连接，并提供广播能力
    """
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


class WebSocketService:
    """
    WebSocket 业务处理，包含接收、存储、回传和广播
    """
    def __init__(self, manager: ConnectionManager):
        self.manager = manager

    async def handle(self, websocket: WebSocket):
        """
        入口方法：接收消息、存储并回传给客户端，同时可选广播给所有人
        """
        # 建立连接
        await self.manager.connect(websocket)
        try:
            while True:
                data = await websocket.receive_text()
                # 存储到数据库
                msg = await Message.create(content=data)
                # 个人回传
                resp = f"[{datetime.now().isoformat()}] 收到：{msg.content}"
                await self.manager.send_personal_message(resp, websocket)
                # （可选）广播给所有其他客户端
                # await self.manager.broadcast(f"用户说：{msg.content}")
        except WebSocketDisconnect:
            self.manager.disconnect(websocket)
